"""
Author: Morphlng
Date: 2024-12-14 20:36:00
LastEditTime: 2024-12-18 11:14:24
LastEditors: Morphlng
Description: 
FilePath: /DrivingGym/src/driving_gym/environment/agent/reward/reward_handler.py
"""

from __future__ import annotations

import os
from typing import TYPE_CHECKING

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.reward.base_reward import BaseReward
from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.misc.util import dynamic_import

if TYPE_CHECKING:
    from driving_gym.simulation.adapter_interface import AdapterInterface


class RewardHandler:
    def __init__(self, config: dict, adapter: "AdapterInterface"):
        self.config = config.copy()
        self._handler = self._create_handler(self.config, adapter)

    def reset(self):
        self._handler.reset()

    def _create_handler(self, config: dict, adapter: "AdapterInterface") -> BaseReward:
        reward_type = config["type"]
        root = os.path.join(os.path.dirname(__file__))
        handler = dynamic_import(
            reward_type,
            root,
            args=(config, adapter),
            suffix="_reward",
        )
        return handler

    def get_reward(
        self, snapshot: Snapshot, scenario: BaseScenario
    ) -> tuple[float, dict]:
        """Get the reward value and details of each reward component.

        Args:
            snapshot (Snapshot): The snapshot containing all data of the current frame.

        Returns:
            reward (float): The reward value.
            details (dict): Details of each reward component.
        """
        return self._handler.get_reward(snapshot, scenario)
